# 安装所需库
# !pip install langchain langchain-experimental langchain-openai pandas matplotlib scikit-learn seaborn

import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from io import StringIO
import tempfile
import base64
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np

from langchain.agents import AgentType, initialize_agent, Tool
from langchain_openai import ChatOpenAI
from langchain.schema import SystemMessage
from langchain.memory import ConversationBufferMemory
from langchain_experimental.agents import create_pandas_dataframe_agent
from langchain.tools import BaseTool
from typing import Type, Optional, Dict, Any
from pydantic import BaseModel, Field

# 设置OpenAI API Key
os.environ["OPENAI_API_KEY"] = "你的OpenAI API Key"


class DataAnalysisAgent:
    def __init__(self, csv_path: str):
        """
        初始化数据探索Agent
        """
        self.csv_path = csv_path
        self.df = pd.read_csv(csv_path)
        self.llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
        self.memory = ConversationBufferMemory(memory_key="chat_history")

        # 创建工具
        self.tools = self._create_tools()

        # 初始化Agent
        self.agent = initialize_agent(
            tools=self.tools,
            llm=self.llm,
            agent=AgentType.OPENAI_FUNCTIONS,
            memory=self.memory,
            verbose=True,
            handle_parsing_errors=True,
            agent_kwargs={
                'system_message': SystemMessage(content=self._get_system_prompt())
            }
        )

    def _get_system_prompt(self) -> str:
        """系统提示词"""
        return """你是一个专业的数据科学家助手。你擅长数据探索、统计分析和机器学习。
        你可以帮助用户：
        1. 提供数据的统计摘要（均值、方差、最大最小值等）
        2. 创建数据可视化图表
        3. 训练机器学习模型并进行预测

        当前数据集信息：
        - 数据集路径: {}
        - 数据形状: {}
        - 列名: {}
        - 前5行数据预览: {}

        请用专业但易懂的方式回答用户问题。如果用户的要求不明确，请询问澄清。
        """.format(
            self.csv_path,
            self.df.shape,
            list(self.df.columns),
            self.df.head().to_string()
        )

    def _create_tools(self) -> list:
        """创建工具列表"""
        return [
            self._create_statistical_summary_tool(),
            self._create_plotting_tool(),
            self._create_ml_training_tool(),
            self._create_data_query_tool()
        ]

    def _create_statistical_summary_tool(self) -> Tool:
        """创建统计摘要工具"""

        class StatisticalSummaryInput(BaseModel):
            columns: Optional[str] = Field(description="需要统计的列名，用逗号分隔。如果为空，则统计所有数值列")

        def statistical_summary(columns: Optional[str] = None) -> str:
            """生成数据的统计摘要"""
            try:
                if columns:
                    column_list = [col.strip() for col in columns.split(',')]
                    # 只选择数值列
                    numeric_columns = self.df[column_list].select_dtypes(include=[np.number]).columns
                    if len(numeric_columns) == 0:
                        return "指定的列中没有数值数据"
                    summary_df = self.df[numeric_columns].describe()
                else:
                    # 统计所有数值列
                    numeric_columns = self.df.select_dtypes(include=[np.number]).columns
                    if len(numeric_columns) == 0:
                        return "数据集中没有数值列"
                    summary_df = self.df[numeric_columns].describe()

                # 添加方差
                variance = self.df[numeric_columns].var()
                summary_df.loc['variance'] = variance

                return f"统计摘要:\n{summary_df.to_string()}"

            except Exception as e:
                return f"生成统计摘要时出错: {str(e)}"

        return Tool(
            name="statistical_summary",
            description="生成数据的统计摘要，包括均值、方差、最大最小值等。输入为需要统计的列名（逗号分隔），如不指定则统计所有数值列。",
            func=statistical_summary,
            args_schema=StatisticalSummaryInput
        )

    def _create_plotting_tool(self) -> Tool:
        """创建绘图工具"""

        class PlottingInput(BaseModel):
            plot_type: str = Field(description="图表类型：histogram, bar, box, scatter, count")
            x_column: str = Field(description="x轴列名")
            y_column: Optional[str] = Field(description="y轴列名（适用于散点图等）")
            title: Optional[str] = Field(description="图表标题")

        def create_plot(plot_type: str, x_column: str, y_column: Optional[str] = None,
                        title: Optional[str] = None) -> str:
            """创建数据可视化图表"""
            try:
                plt.figure(figsize=(10, 6))

                if plot_type == "histogram":
                    if x_column not in self.df.columns:
                        return f"列 '{x_column}' 不存在"
                    plt.hist(self.df[x_column].dropna(), bins=20, alpha=0.7)
                    plt.xlabel(x_column)
                    plt.ylabel('Frequency')

                elif plot_type == "bar":
                    if x_column not in self.df.columns:
                        return f"列 '{x_column}' 不存在"
                    value_counts = self.df[x_column].value_counts()
                    value_counts.plot(kind='bar')
                    plt.xlabel(x_column)
                    plt.ylabel('Count')

                elif plot_type == "box":
                    if x_column not in self.df.columns:
                        return f"列 '{x_column}' 不存在"
                    self.df.boxplot(column=x_column)

                elif plot_type == "scatter":
                    if x_column not in self.df.columns or y_column not in self.df.columns:
                        return f"指定的列不存在"
                    plt.scatter(self.df[x_column], self.df[y_column], alpha=0.6)
                    plt.xlabel(x_column)
                    plt.ylabel(y_column)

                elif plot_type == "count":
                    if x_column not in self.df.columns:
                        return f"列 '{x_column}' 不存在"
                    sns.countplot(data=self.df, x=x_column)

                else:
                    return f"不支持的图表类型: {plot_type}. 支持的类型: histogram, bar, box, scatter, count"

                if title:
                    plt.title(title)
                else:
                    plt.title(f"{plot_type.capitalize()} Plot of {x_column}")

                plt.tight_layout()

                # 保存图片到临时文件并返回base64编码
                with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
                    plt.savefig(tmp_file.name, dpi=300, bbox_inches='tight')
                    with open(tmp_file.name, 'rb') as f:
                        img_base64 = base64.b64encode(f.read()).decode('utf-8')

                plt.close()

                return f"图表已生成！![Plot](data:image/png;base64,{img_base64})"

            except Exception as e:
                return f"创建图表时出错: {str(e)}"

        return Tool(
            name="create_plot",
            description="创建数据可视化图表。支持类型：histogram(直方图), bar(条形图), box(箱线图), scatter(散点图), count(计数图)",
            func=create_plot,
            args_schema=PlottingInput
        )

    def _create_ml_training_tool(self) -> Tool:
        """创建机器学习训练工具"""

        class MLTrainingInput(BaseModel):
            target_column: str = Field(description="目标变量列名")
            feature_columns: str = Field(description="特征列名，用逗号分隔")
            model_type: str = Field(description="模型类型：classification或regression")

        def train_model(target_column: str, feature_columns: str, model_type: str) -> str:
            """训练机器学习模型"""
            try:
                # 检查列是否存在
                if target_column not in self.df.columns:
                    return f"目标列 '{target_column}' 不存在"

                feature_list = [col.strip() for col in feature_columns.split(',')]
                for col in feature_list:
                    if col not in self.df.columns:
                        return f"特征列 '{col}' 不存在"

                # 准备数据
                X = self.df[feature_list]
                y = self.df[target_column]

                # 处理缺失值
                X = X.fillna(X.mean())  # 数值列用均值填充
                y = y.fillna(y.mode()[0] if y.dtype == 'object' else y.mean())  # 分类用众数，数值用均值

                # 分割数据
                X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

                # 训练模型
                if model_type == "classification":
                    model = RandomForestClassifier(n_estimators=100, random_state=42)
                    model.fit(X_train, y_train)
                    y_pred = model.predict(X_test)
                    accuracy = accuracy_score(y_test, y_pred)

                    result = f"""
                    分类模型训练完成！
                    - 目标变量: {target_column}
                    - 特征变量: {', '.join(feature_list)}
                    - 模型类型: 随机森林分类器
                    - 测试集准确率: {accuracy:.4f}
                    - 特征重要性: {dict(zip(feature_list, model.feature_importances_))}
                    """

                elif model_type == "regression":
                    from sklearn.ensemble import RandomForestRegressor
                    from sklearn.metrics import mean_squared_error, r2_score

                    model = RandomForestRegressor(n_estimators=100, random_state=42)
                    model.fit(X_train, y_train)
                    y_pred = model.predict(X_test)
                    mse = mean_squared_error(y_test, y_pred)
                    r2 = r2_score(y_test, y_pred)

                    result = f"""
                    回归模型训练完成！
                    - 目标变量: {target_column}
                    - 特征变量: {', '.join(feature_list)}
                    - 模型类型: 随机森林回归器
                    - 测试集均方误差: {mse:.4f}
                    - 测试集R²分数: {r2:.4f}
                    """
                else:
                    return f"不支持的模型类型: {model_type}. 支持的类型: classification, regression"

                return result

            except Exception as e:
                return f"训练模型时出错: {str(e)}"

        return Tool(
            name="train_machine_learning_model",
            description="训练机器学习模型进行预测。需要指定目标变量、特征变量和模型类型（classification或regression）",
            func=train_model,
            args_schema=MLTrainingInput
        )

    def _create_data_query_tool(self) -> Tool:
        """创建数据查询工具"""

        class DataQueryInput(BaseModel):
            query: str = Field(description="对数据的自然语言查询")

        def query_data(query: str) -> str:
            """回答关于数据的基本问题"""
            try:
                # 使用pandas agent来处理自然语言查询
                pandas_agent = create_pandas_dataframe_agent(
                    self.llm,
                    self.df,
                    verbose=False,
                    allow_dangerous_code=True
                )

                response = pandas_agent.run(query)
                return response

            except Exception as e:
                return f"查询数据时出错: {str(e)}"

        return Tool(
            name="query_data",
            description="回答关于数据的基本问题，如数据基本信息、特定查询等",
            func=query_data,
            args_schema=DataQueryInput
        )

    def chat(self, message: str) -> str:
        """与Agent对话"""
        try:
            response = self.agent.invoke({"input": message})
            return response['output']
        except Exception as e:
            return f"对话过程中出错: {str(e)}"


# 示例使用
def main():
    # 下载Titanic数据集（如果没有的话）
    try:
        df = pd.read_csv("titanic.csv")
    except:
        # 从网络加载Titanic数据集
        url = "https://web.stanford.edu/class/archive/cs/cs109/cs109.1166/stuff/titanic.csv"
        df = pd.read_csv(url)
        df.to_csv("titanic.csv", index=False)

    # 初始化Agent
    agent = DataAnalysisAgent("titanic.csv")

    print("=== 数据探索Agent已就绪 ===")
    print("你可以询问关于数据的问题，请求统计摘要、图表或机器学习模型训练")
    print("输入 'quit' 退出\n")

    # 示例对话
    queries = [
        "请给我数据的统计摘要",
        "创建Survived列的分布图",
        "训练一个预测Survived的分类模型，使用Pclass, Age, Fare作为特征",
        "数据的基本信息是什么？",
        "创建Age的直方图"
    ]

    for i, query in enumerate(queries, 1):
        print(f"\n--- 问题 {i}: {query} ---")
        response = agent.chat(query)
        print(f"回答: {response}")

    # 交互式对话
    print("\n=== 现在开始交互式对话 ===")
    while True:
        user_input = input("\n你的问题 (输入 'quit' 退出): ")
        if user_input.lower() == 'quit':
            break

        response = agent.chat(user_input)
        print(f"\nAgent: {response}")


if __name__ == "__main__":
    main()